import json
from pathlib import Path
from typing import Dict, Optional, Tuple, List

import gradio as gr
import numpy as np
from PIL import Image
import joblib


# 兼容修复：规避 gradio_client 在解析 JSON Schema 布尔 additionalProperties 时的已知问题
# 见报错堆栈 gradio_client.utils._json_schema_to_python_type / get_type 对 bool 处理异常
try:
    from gradio_client import utils as _gcu  # type: ignore

    if hasattr(_gcu, "_json_schema_to_python_type"):
        _orig__json_schema_to_python_type = _gcu._json_schema_to_python_type

        def _patched__json_schema_to_python_type(schema, defs=None):  # type: ignore
            if isinstance(schema, bool):
                return "Any"
            return _orig__json_schema_to_python_type(schema, defs)

        _gcu._json_schema_to_python_type = _patched__json_schema_to_python_type  # type: ignore

    if hasattr(_gcu, "get_type"):
        _orig_get_type = _gcu.get_type

        def _patched_get_type(schema):  # type: ignore
            if isinstance(schema, bool):
                return "Any"
            return _orig_get_type(schema)

        _gcu.get_type = _patched_get_type  # type: ignore
except Exception:
    # 若补丁失败，不影响主流程；仍按原逻辑运行
    pass


ROOT = Path(__file__).resolve().parent
MODEL_PATH = ROOT / "models" / "best_model.pkl"
META_PATH = ROOT / "results" / "best_model.json"


def load_meta() -> Tuple[int, bool]:
    # 默认与训练脚本一致
    img_size = 64
    grayscale = False
    if META_PATH.exists():
        try:
            meta = json.loads(META_PATH.read_text(encoding="utf-8"))
            img_size = int(meta.get("img_size", img_size))
            grayscale = bool(meta.get("grayscale", grayscale))
        except Exception:
            pass
    return img_size, grayscale


def preprocess(img: Image.Image, img_size: int, grayscale: bool) -> np.ndarray:
    if grayscale:
        img = img.convert("L")
    else:
        img = img.convert("RGB")
    img = img.resize((img_size, img_size), Image.BILINEAR)
    arr = np.asarray(img, dtype=np.float32) / 255.0
    arr = arr.reshape(1, -1)
    return arr


class Predictor:
    def __init__(self):
        if not MODEL_PATH.exists():
            raise FileNotFoundError(
                f"未找到模型: {MODEL_PATH}. 请先运行 train.py 训练并保存最佳模型。"
            )
        self.model = joblib.load(MODEL_PATH)
        self.img_size, self.grayscale = load_meta()

        # 尝试获取类名
        est = self.model.named_steps.get("estimator", None)
        self.classes_ = None
        if est is not None and hasattr(est, "classes_"):
            try:
                self.classes_ = list(est.classes_)
            except Exception:
                pass

    def predict(self, img: Image.Image) -> Dict[str, float]:
        x = preprocess(img, self.img_size, self.grayscale)
        # 优先使用 predict_proba
        try:
            proba = self.model.predict_proba(x)
            if self.classes_ is None and hasattr(self.model.named_steps.get("estimator", None), "classes_"):
                try:
                    self.classes_ = list(self.model.named_steps["estimator"].classes_)
                except Exception:
                    pass
            if self.classes_ is None:
                # 没有类名则用索引
                self.classes_ = [str(i) for i in range(proba.shape[1])]
            return {cls: float(prob) for cls, prob in zip(self.classes_, proba[0].tolist())}
        except Exception:
            # 回退：无 predict_proba 时用 predict 给出 1/0 置信
            pred = self.model.predict(x)[0]
            if self.classes_ is None:
                return {str(pred): 1.0}
            return {cls: (1.0 if cls == pred else 0.0) for cls in self.classes_}


predictor: Optional[Predictor] = None


def infer(img: Image.Image) -> Tuple[str, List[List]]:
    global predictor
    if predictor is None:
        predictor = Predictor()
    probs = predictor.predict(img)
    # 选取 Top-1 预测
    if probs:
        top1 = max(probs.items(), key=lambda x: x[1])[0]
    else:
        top1 = "(no-prediction)"
    # 转为表格 [[class, prob], ...]，按概率降序
    table = [[k, float(v)] for k, v in sorted(probs.items(), key=lambda x: x[1], reverse=True)]
    return top1, table


def build_interface() -> gr.Interface:
    title = "Cats vs Dogs Classifier (Best Model)"
    description = "上传一张图片，使用训练效果最好的模型进行猫狗分类。"
    examples = []
    # 从 data/test 里收集少量 cat/dog 示例
    for sub in [ROOT / "data" / "test", ROOT / "data" / "train"]:
        if sub.exists():
            for cls in ["cat", "dog"]:
                for p in (sub / cls).glob("*.*"):
                    examples.append(str(p))
                    if len(examples) >= 6:
                        break
                if len(examples) >= 6:
                    break
        if len(examples) >= 6:
            break

    return gr.Interface(
        fn=infer,
        inputs=gr.Image(type="pil", label="上传猫或狗的图片"),
        outputs=[
            gr.Textbox(label="预测类别"),
            gr.Dataframe(headers=["类别", "概率"], datatype=["str", "number"], label="预测概率分布"),
        ],
        title=title,
        description=description,
        examples=examples if examples else None,
        allow_flagging="never",
    )


if __name__ == "__main__":
    app = build_interface()
    # 本地启动，默认 7860 端口；可通过 --server.port 指定
    # 避免某些环境下无法访问 localhost 的报错，默认开启 share
    # 不固定端口，让 Gradio 自动选择空闲端口
    app.launch(server_name="0.0.0.0", server_port=None, share=True)
